# -*- coding: utf-8 -*-

import torch
import segmentation_models_pytorch as smp
from CC_CCI_dataset import get_TCSM_dataloader
from model import TCSM
from losses import TCSMLoss
from trainer import TCSMEpoch
from metrics import Accuracy, IoU

IMAGE_DIR = "D:\data\ct_lesion_seg\image"
MASK_DIR = "D:\data\ct_lesion_seg\mask"
CLASSES = 3
BATCH_SIZE = 5
LABELED_RATE = 0.5
times = int(1 / LABELED_RATE - 1)


train_loader, valid_loader, unlabeled_loader = get_TCSM_dataloader(IMAGE_DIR, MASK_DIR, 
                                                                   labeled_batch_size=BATCH_SIZE,
                                                                   unlabeled_times=times,
                                                                   device='cuda:0')

student = smp.DeepLabV3Plus(encoder_name="mobilenet_v2", encoder_weights=None,
                            in_channels=1, classes=CLASSES, activation='softmax2d')
teacher = smp.DeepLabV3Plus(encoder_name="mobilenet_v2", encoder_weights=None,
                            in_channels=1, classes=CLASSES, activation='softmax2d')
tcsm = TCSM(student, teacher, student_device='cuda:0', teacher_device='cpu')

loss = TCSMLoss(device='cuda:0')
optimizer = torch.optim.Adam(params=tcsm.parameters())
metrics = [
    Accuracy(),
    IoU()
]

tcsmEpoch = TCSMEpoch(tcsm, loss=loss, optimizer=optimizer, metrics=metrics, verbose = True)

for epoch in range(1):
    print("epoch: {}".format(epoch))
    tcsmEpoch.train(train_loader, unlabeled_loader, epoch)
    tcsmEpoch.valid(valid_loader)